526D - Om Nom and Necklace - CodeForces Solution

#include <bits/stdc++.h>

using namespace std;

#pragma GCC target("avx2")

#ifdef LOCAL
#include "algo/debug.h"
#endif

template<class Fun>
class y_combinator_result {
   Fun fun_;
public:
   template<class T>
   explicit y_combinator_result(T &&fun): fun_(std::forward<T>(fun)) {}

   template<class ...Args>
   decltype(auto) operator()(Args &&...args) {
      return fun_(std::ref(*this), std::forward<Args>(args)...);
   }
};

template<class Fun>
decltype(auto) y_combinator(Fun &&fun) {
   return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun));
}

namespace str {
/** Computes the Pi array of s. */
vector<int> kmp(const string &s) {
   const int n = (int) s.size();
   vector<int> pi(n);
   for (int i = 1, j = 0; i < n; ++i) {
      while (j > 0 && s[j] != s[i]) {
         j = pi[j - 1];
      }
      j += (int) (s[j] == s[i]);
      pi[i] = j;
   }
   return pi;
}
}

int main() {
   ios::sync_with_stdio(false);
   cin.tie(nullptr);
   int n, k;
   cin >> n >> k;
   string s;
   cin >> s;
   vector<int> pi = str::kmp(s);
   vector<vector<int>> tree(n + 1);
   for (int i = 0; i < n; ++i) {
      tree[pi[i]].push_back(i + 1);
   }
   vector<bool> visited(n + 1);
   vector<int> enter(n + 1), exit(n + 1);
   int timer = 0;
   auto dfs = y_combinator([&](auto self, int node) -> void {
      visited[node] = true;
      enter[node] = timer++;
      for (int child : tree[node]) {
         if (!visited[child]) {
            self(child);
         }
      }
      exit[node] = timer++;
   });
   dfs(0);
   auto ancesotr = [&](int a, int b) -> bool {
      return enter[a] <= enter[b] && exit[b] <= exit[a];
   };
   vector<bool> possible(n + 1, true);
   for (int i = 1; i <= n; ++i) {
      for (int j = 1; j <= min(k, n / i); ++j) {
         if (!ancesotr(i, j * i)) {
            possible[i] = false;
            break;
         }
      }
   }
   visited.assign(n + 1, false);
   vector<bool> ans(n + 1);
   vector<set<int>> track(k);
   auto solve = y_combinator([&](auto self, int node) -> void {
      visited[node] = true;
      int i = node, j = i / (k + 1);
      if (j * (k + 1) == i && possible[j] && ancesotr(j, i)) {
         ans[i] = true;
      }
      else if (!track[i % k].empty()) {
         auto it = track[i % k].lower_bound((i + k) / (k + 1));
         while (it != track[i % k].begin()) {
            it = prev(it);
            j = *it;
            if (possible[(i - j) / k]) {
               ans[i] = true;
               break;
            }
         }
      }
      track[i % k].insert(i);
      for (int child : tree[node]) {
         if (!visited[child]) {
            self(child);
         }
      }
      track[i % k].erase(i);
   });
   solve(0);
   for (int i = 1; i <= n; ++i) {
      putchar(ans[i] ? '1' : '0');
   }
}

/* stuff you should look for
 * int overflow, array bounds
 * special cases (n=1?)
 * do smth instead of nothing and stay organized
 * WRITE STUFF DOWN
 * DON'T GET STUCK ON ONE APPROACH
 */